import json
import random
import warnings

import networkx as nx
import matplotlib.pyplot as plt
from typing import Dict, List, Optional
import os
from ipaddress import IPv4Network, IPv4Address

import numpy as np


class TopologyGenerator:
    def __init__(self):
        self.switches = []  # 所有交换机列表
        self.switch_types = {
            "switchA": [],  # EGRESS switches
            "switchB": [],  # INGRESS switches
            "switchC": []  # NORMAL switches
        }
        self.links = []
        self.applications = []
        self.cpn_nodes = []

        # IP地址管理
        self.used_ips = set()
        self.ip_network = IPv4Network('10.0.0.0/16')
        self.anycast_network = IPv4Network('192.168.255.0/24')

        # 基础配置
        self.topology_params = {
            "random_seed": 42,
            "cpu_period": 100000,
            "max_cpu_util": 0.8,
            "min_cpu_period": 1000,
            "config_host_path": "/home/u/ryu/ryu/app/SA-CPN/config",
        }

        # 服务类型配置
        self.service_configs = self._generate_service_configs()

    def _generate_service_configs(self):
        """生成服务配置，确保anycast IP不重复"""
        services = {
            "COMPUTE": {
                "max_delay_range": (20, 5),
                "compute_load_range": (100,20),
                "data_size_range": (100, 20),
                "description": "compute-intensive-app"
            },
            "IO": {
                "max_delay_range": (20, 5),
                "compute_load_range":(20, 10),
                "data_size_range": (1000, 20),
                "description": "io-intensive-app"
            },
            "HYBRID": {
                "max_delay_range": (20, 5),
                "compute_load_range": (20, 10),
                "data_size_range": (100, 20),
                "description": "hybrid-app"
            },
            "LATENCY": {
                "max_delay_range": (5, 1),
                "compute_load_range": (20, 10),
                "data_size_range": (100, 20),
                "description": "latency-sensitive-app"
            }
        }

        # 为每个服务分配anycast IP
        anycast_ips = list(self.anycast_network.hosts())
        for i, (service_type, config) in enumerate(services.items()):
            config["anycast_ip"] = str(anycast_ips[i])
            config["port"] = 8000
            config["target_nodes"] = config["anycast_ip"] + ":" + str(config["port"])

        return services

    def _get_next_ip(self) -> str:
        """获取下一个可用的IP地址"""
        for ip in self.ip_network.hosts():
            if str(ip) not in self.used_ips:
                self.used_ips.add(str(ip))
                return str(ip)
        raise Exception("IP地址耗尽")

    def _generate_wan_links(self):
        """生成广域网链路"""
        # EGRESS连接到NORMAL
        for egress in self.switches:
            if egress["type"] == "EGRESS":
                normal = random.choice(self.switches)
                while normal["type"] != "NORMAL":
                    normal = random.choice(self.switches)
                self.links.append(self._create_link(egress["name"], normal["name"]))

        # INGRESS连接到NORMAL
        for ingress in self.switches:
            if ingress["type"] == "INGRESS":
                normal = random.choice(self.switches)
                while normal["type"] != "NORMAL":
                    normal = random.choice(self.switches)
                self.links.append(self._create_link(ingress["name"], normal["name"]))

        # NORMAL之间互连
        for i, switch1 in enumerate(self.switches):
            for switch2 in self.switches[i + 1:]:
                if random.random() < 0.5:  # 50%概率连接
                    self.links.append(self._create_link(switch1["name"], switch2["name"]))

    def _create_link(self, source: str, target: str):
        """创建链路"""
        return {
            "source": source,
            "target": target,
            "bw": random.randint(5, 20),  # 带宽5-20Mbps
            "delay": f"{random.randint(5, 30)}ms",  # 延迟1-10ms
            "jitter": f"{random.randint(1, 2)}ms",
            "loss":0 # round(random.uniform(0.01, 0.2), 3)  # 丢包率0.1%-20%
        }

    def import_from_zoo(self, filename: str):
        """从TopologyZoo导入拓扑"""
        G = nx.read_graphml(filename)

        # 转换节点
        self.switches = []
        for i, node in enumerate(G.nodes()):
            self.switches.append({
                "name": f"s{i + 1}",
                "dpid": i + 1,
                "type": "NORMAL",  # 初始都标记为NORMAL
            })

        # 转换链路
        self.links = []
        for edge in G.edges():
            self.links.append(self._create_link(
                f"s{list(G.nodes()).index(edge[0]) + 1}",
                f"s{list(G.nodes()).index(edge[1]) + 1}"
            ))

    def distribute_users(self, user_count: int):
        """分布用户节点并标记INGRESS交换机"""
        self.cpn_nodes = []
        available_switches = [s for s in self.switches]
        base_port = 6000
        for i in range(user_count):
            # 随机选择一个交换机连接用户
            switch = random.choice(available_switches)
            switch["type"] = "INGRESS"  # 标记为INGRESS
            self.switch_types["switchB"].append(switch)

            # 创建用户节点
            self.cpn_nodes.append({
                "name": f"cpNode{i + 1}",
                "ip": self._get_next_ip(),
                "port": base_port,
                "port_binding": base_port + len(self.cpn_nodes),
                "connected_to": switch["name"],
                "capacity": 50000,
                "bw": 10
            })

    def distribute_applications(self, app_counts: Dict[str, int]):
        """分布应用实例并标记EGRESS交换机"""
        self.applications = []
        available_switches = [s for s in self.switches]

        base_port = 7000
        for app_type, count in app_counts.items():
            for i in range(count):
                # 随机选择一个交换机连接应用
                switch = random.choice(available_switches)
                switch["type"] = "EGRESS"  # 标记为EGRESS
                self.switch_types["switchA"].append(switch)

                # 创建应用实例
                self.applications.append({
                    "name": f"{app_type.lower()}{i + 1}",
                    "type": app_type,
                    "ip": self._get_next_ip(),
                    "port": base_port,
                    "port_binding": base_port + len(self.applications),
                    "connected_to": switch["name"],
                    "capacity": 50000,
                    "request_handle_capacity": random.randint(40, 60)
                })

    def classify_switches(self):
        """分类剩余的交换机为NORMAL"""
        self.switch_types["switchC"] = [
            s for s in self.switches
            if s["type"] == "NORMAL"
        ]

    def save_topology(self, filename: str):
        """保存拓扑配置到文件"""
        topology = {
            "topology_params": self.topology_params,
            "switches": self.switch_types,
            "all_switches": self.switches,
            "links": self.links,
            "service_configs": self.service_configs,
            "applications": self.applications,
            "cpn_nodes": self.cpn_nodes
        }

        with open(filename, 'w') as f:
            json.dump(topology, f, indent=4)

    def visualize_topology(self, filename: str):
        """修复边显示问题的完整可视化代码"""
        plt.close('all')
        G = nx.Graph()
        pos = {}
        node_colors = []
        node_sizes = []
        labels = {}
        edge_labels = {}

        # 构建基础网络拓扑
        switch_names = [s["name"] for s in self.switches]
        temp_G = nx.Graph()
        temp_G.add_nodes_from(switch_names)

        # 添加所有链路（核心修正点）
        all_edges = []
        for link in self.links:
            temp_G.add_edge(link["source"], link["target"])
            all_edges.append((link["source"], link["target"]))

        # 混合布局算法增强
        base_pos = nx.spring_layout(temp_G,
                                    k=0.6,  # 增大节点间距
                                    iterations=500,  # 保证布局收敛
                                    seed=42,
                                    scale=3.0)

        # 添加交换机节点
        for switch in self.switches:
            G.add_node(switch["name"])
            pos[switch["name"]] = base_pos[switch["name"]]
            labels[switch["name"]] = f"{switch['name']}\n({switch['type']})"

            # 根据类型设置样式
            if switch["type"] == "EGRESS":
                node_colors.append('#1f78b4')  # 深蓝
                node_sizes.append(1500)
            elif switch["type"] == "INGRESS":
                node_colors.append('#e31a1c')  # 红色
                node_sizes.append(1300)
            else:
                node_colors.append('#33a02c')  # 绿色
                node_sizes.append(1100)

        # 用户节点布局（左侧扇形）
        user_radius = 0.8
        user_angle = np.linspace(120, 60, len(self.cpn_nodes))  # 120°到60°
        for idx, user in enumerate(self.cpn_nodes):
            G.add_node(user["name"])
            switch_pos = pos[user["connected_to"]]

            # 极坐标转换
            theta = np.radians(user_angle[idx])
            dx = user_radius * np.cos(theta)
            dy = user_radius * np.sin(theta)
            pos[user["name"]] = (switch_pos[0] + dx, switch_pos[1] + dy)

            labels[user["name"]] = f"{user['ip']}\n:{user['port']}"
            node_colors.append('#ff7f00')  # 橙色
            node_sizes.append(700)
            G.add_edge(user["name"], user["connected_to"])

        # 应用节点布局（右侧扇形）
        app_radius = 0.8
        app_angle = np.linspace(-60, -120, len(self.applications))  # -60°到-120°
        for idx, app in enumerate(self.applications):
            G.add_node(app["name"])
            switch_pos = pos[app["connected_to"]]

            theta = np.radians(app_angle[idx])
            dx = app_radius * np.cos(theta)
            dy = app_radius * np.sin(theta)
            pos[app["name"]] = (switch_pos[0] + dx, switch_pos[1] + dy)

            labels[app["name"]] = f"{app['ip']}\n:{app['port']}"
            node_colors.append('#6a3d9a')  # 紫色
            node_sizes.append(700)
            G.add_edge(app["name"], app["connected_to"])

        # 关键修复1：显式添加交换机间边
        for u, v in all_edges:
            G.add_edge(u, v)

        # 关键修复2：动态斥力调整
        fixed_switches = {s["name"]: pos[s["name"]] for s in self.switches}
        pos = nx.spring_layout(G, pos=pos,
                               fixed=fixed_switches.keys(),
                               k=0.2,  # 外围节点斥力
                               iterations=100,
                               seed=42)

        # 可视化参数设置
        plt.figure(figsize=(36, 24), dpi=300)

        # 绘制核心交换机边（关键修复3）
        core_edges = [(u, v) for u, v in G.edges()
                      if u in switch_names and v in switch_names]
        nx.draw_networkx_edges(
            G, pos,
            edgelist=core_edges,
            width=3.5,  # 加宽核心边
            alpha=0.9,
            edge_color='#2b2d42',  # 深蓝色
            style='solid',
            arrows=False
        )

        # 绘制接入边
        access_edges = [e for e in G.edges() if e not in core_edges]
        nx.draw_networkx_edges(
            G, pos,
            edgelist=access_edges,
            width=1.2,
            alpha=0.6,
            edge_color='#7f7f7f',  # 灰色
            style='dashed'
        )

        # 绘制节点（分层次绘制）
        nx.draw_networkx_nodes(
            G, pos,
            nodelist=switch_names,
            node_color=node_colors[:len(switch_names)],
            node_size=node_sizes[:len(switch_names)],
            edgecolors='black',
            linewidths=2.0
        )
        nx.draw_networkx_nodes(
            G, pos,
            nodelist=[n["name"] for n in self.cpn_nodes + self.applications],
            node_color=node_colors[len(switch_names):],
            node_size=node_sizes[len(switch_names):],
            edgecolors='black',
            linewidths=1.0
        )

        # 标签优化
        text_pos = {k: (v[0], v[1] + 0.05) for k, v in pos.items()}
        nx.draw_networkx_labels(
            G, text_pos,
            labels=labels,
            font_size=10,
            font_family='DejaVu Sans',
            verticalalignment='bottom',
            bbox=dict(
                # facecolor='white',
                alpha=0.85,
                edgecolor='none',
                boxstyle='round,pad=0.3'
            )
        )

        # 专业图例
        legend_elements = [
            plt.Line2D([0], [0], marker='o', color='w', markersize=20,
                       markerfacecolor='#1f78b4', label='EGRESS Switch'),
            plt.Line2D([0], [0], marker='o', color='w', markersize=20,
                       markerfacecolor='#e31a1c', label='INGRESS Switch'),
            plt.Line2D([0], [0], marker='o', color='w', markersize=16,
                       markerfacecolor='#33a02c', label='CORE Switch'),
            plt.Line2D([0], [0], marker='o', color='w', markersize=14,
                       markerfacecolor='#ff7f00', label='User Nodes'),
            plt.Line2D([0], [0], marker='o', color='w', markersize=14,
                       markerfacecolor='#6a3d9a', label='Applications'),
            plt.Line2D([0], [0], color='#2b2d42', lw=4, label='Core Links'),
            plt.Line2D([0], [0], color='#7f7f7f', lw=2, linestyle='dashed', label='Access Links')
        ]

        plt.legend(
            handles=legend_elements,
            loc='upper center',
            bbox_to_anchor=(0.5, 1.08),
            ncol=4,
            fontsize=14,
            frameon=False
        )

        plt.title("Optimized Network Topology with Visible Edges", fontsize=20, pad=25)
        plt.axis('off')
        plt.savefig(filename, bbox_inches='tight', dpi=300)
        plt.close()

def main():
    generator = TopologyGenerator()

    # 1. 导入基础拓扑
    generator.import_from_zoo("Abilene.graphml")  # 使用实际的广域网拓扑

    app_counts = {
        "COMPUTE": 10,
        "IO": 3,
        "HYBRID": 8,
        "LATENCY": 5
    }
    # 3. 分布应用实例并标记EGRESS交换机
    generator.distribute_applications(app_counts)

    # 2. 分布用户并标记INGRESS交换机
    generator.distribute_users(user_count=20)

    # 4. 分类剩余交换机为NORMAL
    generator.classify_switches()

    # 5. 保存配置和可视化
    generator.save_topology("topo_gen/topo_config.json")
    generator.visualize_topology("topo_gen/topology.png")


if __name__ == "__main__":
    main() 